import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
# from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
import pymysql
# from sklearn.metrics import r2_score, mean_absolute_error

db_conf = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
}
engire = create_engine(
    f"mysql+pymysql://{db_conf['user']}:{db_conf['password']}@{db_conf['host']}:{db_conf['port']}/{db_conf['database']}")
conn = pymysql.connect(**db_conf)
chunk_size = 10000
df = pd.read_sql_query("""
    SELECT d.* FROM date_1 d WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
    """,
    conn,
    chunksize=chunk_size)
df1 = pd.concat(df, ignore_index=True)

df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
print(df1.head)

# 处理缺少数据
df1 = df1.dropna(subset=['zd_closes'])
print(df1.head)
ex = ['id', 'ts_code', 'trade_date','the_date','opens','high','low','closes','pre_closes','changes','pct_chg', 'amount']

# 筛选合适的自变量
number = df1.select_dtypes(include=['number']).columns.tolist()
newList  = [col for col in number if col not in ex]

# print(newList)
formuls = 'zd_closes ~ ' + ' + '.join(newList)
res = smf.ols(formuls, data=df1).fit()
print(res.summary())
